package com.hle.map;
import java.util.Comparator;
import java.util.LinkedList;
import java.util.Queue;


public class TreeMap<K, V> implements Map<K, V> {
	private static final boolean RED = false;
	private static final boolean BLACK = true;
	
	private int size;
	private Node<K, V> root;
	private Comparator<K> comparator;
	
	public TreeMap() {
		this(null);
	}
	
	public TreeMap(Comparator<K> comparator) {
		this.comparator = comparator;
	}
	
	public int size() {
		return size;
	}

	public boolean isEmpty() {
		return size == 0;
	}

	public void clear() {
		root = null;
		size = 0;
	}

	@Override
	public V put(K key, V value) {
		// TODO Auto-generated method stub
		elementNotNullCheck(key);
		if (root == null) {
			root = new Node<>(key,value,null);
			size++;
			afterPut(root);
			return null;
		}
		Node<K,V> node = root;
		Node<K,V> parent = null;

		int cmp = 0;
		while (node != null) {
			cmp = compare(key, node.key);
			parent = node;
			if (cmp > 0) {
				node = node.right;
			} else if (cmp < 0) {
				node = node.left;
			} else {
				V oldValue = node.value;
				node.key = key;
				node.value = value;
				return oldValue;
			}
		}
		Node<K,V> newNode = new Node<>(key, value, parent);
		
		if (cmp > 0) {
			parent.right = newNode;
		} else {
			parent.left = newNode;
		}
		size++;	
		afterPut(newNode);
		return null;
	}

	@Override
	public V get(K key) {
		Node<K, V> node = node(key);
		return node == null ? null : node.value;
	}

	@Override
	public V remove(K key) {
		Node<K, V> node = node(key);
		if (node == null) return null;
		return remove(node);
	}

	@Override
	public boolean containsKey(K key) {
		Node<K, V> node = node(key);
		return node == null ? false : true;
	}

	@Override
	public boolean containsValue(V value) {
		// TODO Auto-generated method stub
		if (root == null) return false;
		Queue<Node<K, V>> queue = new LinkedList<>();
		queue.offer(root);
		while(!queue.isEmpty()) {
			Node<K, V> node = queue.poll();
			if (valEquals(node.value, value))return true;
			if (node.left != null) queue.offer(node.left);
			if (node.right != null) queue.offer(node.right);
		}
		return false;
	}

	@Override
	public void traversal(Visitor<K, V> visitor) {
		// TODO Auto-generated method stub
		if (root == null) return;
		traversal(root, visitor);
	}
	
	public void traversal(Node<K, V> node,Visitor<K, V> visitor) {
		 if (visitor == null || node == null) return;
		traversal(node.left, visitor);
		boolean stop = visitor.visit(node.key, node.value);
		if (stop) return;
		traversal(node.right, visitor);
	}
	
    private	void elementNotNullCheck(K key) {
		if (key == null) {
			throw new IllegalArgumentException("key must not be null");
		}
	}
    
    private void afterPut(Node<K, V> node) {
    	Node<K, V> parent  = node.parent;
		if (parent == null) {
			color(node, BLACK);
			return;
		}
		
		if (isBlack(parent)) {
			return;
		}
		
		Node<K, V> uncle = parent.sibling();
		Node<K, V> grand = parent.parent;
		
		if (isRed(uncle)) {
			black(parent);
			black(uncle);
			red(grand);
			afterPut(grand);
			return;
		}else {
			if (parent.isLeftChild()) { // L
				if (node.isLeftChild()) { // LL
					black(parent);
					red(grand);
					rotateRight(grand);
				}else { //LR
					black(node);
					red(grand);
					rotateLeft(parent);
					rotateRight(grand);
				}
			}else { // R
				if (node.isLeftChild()) { //RL
					black(node);
					red(grand);
					rotateRight(parent);
					rotateLeft(grand);
				}else { //RR
					black(parent);
					red(grand);
					rotateLeft(grand);
				}
			}
		}	
	}
    
    public void afterRemove(Node<K,V> node, Node<K,V> replacement) {
		// TODO Auto-generated method stub
		if (isRed(node)) return;
		
		if (isRed(replacement)) {
			black(replacement);
			return;
		}
		Node<K,V> parent = node.parent;
		if (parent == null) return;
		
		///  删除的是黑色叶子节点
		boolean left = parent.left == null || node.isLeftChild();
		Node<K,V> sibling = left ? parent.right : parent.left;
			
		if (left) {
			if (isRed(sibling)) {
				black(sibling);
				red(parent);
				rotateLeft(parent);
				sibling = parent.right;
			}
			
			/// 兄弟节点是黑色
			if (isBlack(sibling.left) && isBlack(sibling.right)) {
				Boolean parentBlack = isBlack(parent);
				red(sibling);
				black(parent);
				if (parentBlack) {
					afterRemove(parent, null);
				}
			}else {
				if(isBlack(sibling.left) && isRed(sibling.right)) { // 黑兄弟左边没有节点
					rotateRight(sibling);
					sibling = parent.right;
				}
				color(sibling, colorOf(parent));
				black(parent);
				black(sibling.right);
				rotateLeft(parent);
			}
		}else {
			if (isRed(sibling)) {
				black(sibling);
				red(parent);
				rotateRight(parent);
				sibling = parent.left;
			}
			
			/// 兄弟节点是黑色
			if (isBlack(sibling.left) && isBlack(sibling.right)) {
				Boolean parentBlack = isBlack(parent);
				red(sibling);
				black(parent);
				if (parentBlack) {
					afterRemove(parent, null);
				}
			}else {
				if(isBlack(sibling.left) && isRed(sibling.right)) { // 黑兄弟左边没有节点
					rotateLeft(sibling);
					sibling = parent.left;
				}
				color(sibling, colorOf(parent));
				black(parent);
				black(sibling.left);
				rotateRight(parent);
			}
			
		}
		
	}	
    
    private boolean isBlack(Node<K, V> node) {
		return colorOf(node) == BLACK;
	}
	
	private boolean isRed(Node<K, V> node) {
		return colorOf(node) == RED;
	}
	
	private boolean colorOf(Node<K, V> node) {
		return node == null ? BLACK : node.color;
	}
	
	private Node<K, V> color(Node<K, V> node, boolean color) {
		if (node != null) node.color = color;
		return node;
	}
	
	
	private Node<K, V> black(Node<K, V> node) {
		return color(node, BLACK);
	}
	
	private Node<K, V> red(Node<K, V> node) {
		return color(node, RED);
	}
	
	 public void rotateLeft(Node<K,V> node) {
			Node<K,V> parent = node.right;
			Node<K,V> child = parent.left;
			node.right = child;
			parent.left = node;
			afterRotate(node, parent, null);
		}
		
	    public void rotateRight(Node<K,V> node) {
			Node<K,V> parent = node.left;
			Node<K,V> child = parent.right;
			node.left = child;
			parent.right = node;
			afterRotate(node, parent, null);
		}
		
		public void afterRotate(Node<K,V> grand, Node<K,V> parent, Node<K,V> child) {
			// 让parent称为子树的根节点
			parent.parent = grand.parent;
			if (grand.isLeftChild()) {
				grand.parent.left = parent;
			} else if (grand.isRightChild()) {
				grand.parent.right = parent;
			} else { // grand是root节点
				root = parent;
			}
			
			// 更新child的parent
			if (child != null) {
				child.parent = grand;
			}
			
			// 更新grand的parent
			grand.parent = parent;
		}
		
		public void rotate(
				Node<K,V> r, // 子树根节点
				Node<K,V> a, Node<K,V> b, Node<K,V> c,
				Node<K,V> d,
				Node<K,V> e, Node<K,V>f, Node<K,V>g) {
			
			d.parent = r.parent;
			if (r.isLeftChild()) {
				r.parent.left = d;
			}else if (r.isRightChild()) {
				r.parent.right = d;
			}else {
				root = d;
			}
			
			/// a - b - c
			
			b.left = a;
			if (a != null) {
				a.parent = b;
			}
			b.right = c;
			if (c != null) {
				c.parent = c;
			}
			
			
			
			/// e - f - g
			f.left = e;
			if (e != null) {
				e.parent = f;
			}
			b.right = g;
			if (g != null) {
				g.parent = f;
			}
			
			// b - d - f
			
			d.left = b;
			d.right = f;
			b.parent = d;
			f.parent = d;
		}
    
	@SuppressWarnings("unchecked")
	private int compare(K e1, K e2) {
		if (comparator == null) {
			return ((Comparable<K>) e1).compareTo(e2);
		}
		return comparator.compare(e1, e2);
	}
	
	private Node<K,V> node(K key) {
		Node<K,V> node = root;
		while (node != null) {
			int cmp = compare(key, node.key);
			if (cmp == 0) {
				return node;
			}else if (cmp > 0) {
				node = node.right;
			}else {
				node = node.left;
			}
		}
		return null;
	}
    
	public V remove(Node<K, V> node) {
		if (node == null) return null;
		size--;
		
		if (node.hasTwoChildren()) { // 度为2的节点
			// 找到后继节点
			Node<K,V> s = successor(node);
			// 用后继节点的值覆盖度为2的节点的值
			node.key = s.key;
			node.value = s.value;
			// 删除后继节点
			node = s;
		}
		
		// 删除node节点（node的度必然是1或者0）
		Node<K,V> replacement = node.left != null ? node.left : node.right;
		
		if (replacement != null) { // node是度为1的节点
			// 更改parent
			replacement.parent = node.parent;
			// 更改parent的left、right的指向
			if (node.parent == null) { // node是度为1的节点并且是根节点
				root = replacement;
			} else if (node == node.parent.left) {
				node.parent.left = replacement;
			} else { // node == node.parent.right
				node.parent.right = replacement;
			}
			
			// 删除节点之后的处理
			afterRemove(node,replacement);
		} else if (node.parent == null) { // node是叶子节点并且是根节点
			root = null;
			
			// 删除节点之后的处理
			afterRemove(node,null);
		} else { // node是叶子节点，但不是根节点
			if (node == node.parent.left) {
				node.parent.left = null;
			} else { // node == node.parent.right
				node.parent.right = null;
			}
			
			// 删除节点之后的处理
			afterRemove(node,null);
		}
		return null;
	}
	

    private Node<K,V> successor(Node<K,V> node) {
		if (node == null) {
			return null;
		}
		Node<K,V> p = node.left;
		if (p != null) {
			while (p.right != null) {
				p= p.right;
			}
			return p;
		}
		
		while (node.parent != null && node == node.parent.right) {
			node = node.parent;
		}
		
		return node.parent;
	}
    	
    private boolean valEquals(V v1, V v2) {
    	return v1 == null ? v2 == null : v1.equals(v2);
	}
	
	private static class Node<K, V> {
		K key;
		V value;
		boolean color = RED;
		Node<K, V> left;
		Node<K, V> right;
		Node<K, V> parent;
		public Node(K key, V value, Node<K, V> parent) {
			this.key = key;
			this.value = value;
			this.parent = parent;
		}
		
		public boolean isLeaf() {
			return left == null && right == null;
		}
		
		public boolean hasTwoChildren() {
			return left != null && right != null;
		}
		
		public boolean isLeftChild() {
			return parent != null && this == parent.left;
		}
		
		public boolean isRightChild() {
			return parent != null && this == parent.right;
		}
		
		public Node<K, V> sibling() {
			if (isLeftChild()) {
				return parent.right;
			}
			
			if (isRightChild()) {
				return parent.left;
			}
			return null;
		}
	}
	
}
